# Tools/database.py

from sqlalchemy import create_engine, text
from sqlalchemy.engine.base import Engine
from langchain_community.utilities import SQLDatabase
from typing import Optional, Dict, Any


class UniversalDBConnector:

    def __init__(self, connection_str: str):
        """
        :param connection_str: 数据库连接字符串
           示例:
           - SQLite: sqlite:///example.db
           - MySQL: mysql+pymysql://user:pass@host/dbname
           - PostgreSQL: postgresql://user:pass@host/dbname
        """
        self.engine = create_engine(connection_str)
        self.db = SQLDatabase(self.engine)

    def get_table_info(self) -> str:
        """获取数据库表结构信息用于提示工程"""
        return self.db.get_table_info()

    def execute_safe_query(self, query: str):
        """
        执行SQL语句
        :param query: SQL语句
        :return: 执行结果
        """
        if any(cmd in query.upper() for cmd in ["ALTER"]):
            raise ValueError("写操作被禁止")

        # 处理写操作
        if any(cmd in query.upper() for cmd in ["DELETE", "UPDATE", "INSERT"]):
            with self.engine.connect() as conn:
                try:
                    # 这里修改：执行操作并返回结果对象
                    result = conn.execute(text(query))
                    conn.commit()
                    return result  # 返回执行结果对象而不是连接对象
                except Exception as e:
                    conn.rollback()
                    raise e

        # 处理查询操作
        limited_query = query
        with self.engine.connect() as conn:
            result = conn.execute(text(limited_query))
            return result

    @property
    def dialect(self) -> str:
        """获取数据库方言"""
        return self.db.dialect
